# -*- coding: utf-8 -*-
# created on 2016/12/5

from sympy.abc import n
from sympy import simplify, solve, sympify
import re


# 设置默认变量为 n
def sl_default_symbol(expr, args=None):
    if args is None:
        args = []
    if len(args) > 1:
        return args[1].toSympy()
    elif n in expr.free_symbols:
        return n
    else:
        return expr.free_symbols.pop()


# 求递推公式的项: 含*_(*)
def sl_xiangs(poly):
    p = r'[a-zA-Z]_\(.*?\)'
    m = re.findall(p, str(poly))
    m = list(set(m))
    return m


# 求递推公式的项: 含*_(?n)
def sl_xiangs001(poly):
    p = r'[a-zA-Z]_\(.*?\)'
    m = re.findall(p, str(poly))
    m = list(set(m))
    xiangs = []
    for item in m:
        new_item = sympify(item)
        num = new_item.args[0]
        if num.has(n):
            xiangs.append(item)
    return xiangs


# 求递推公式下标的最值
def sl_max_min_index(poly):
    p = r'[a-zA-Z]_\(.*?\)'
    m = re.findall(p, str(poly))
    m = list(set(m))
    m_indexs = []
    for item in m:
        new_item = simplify(item)
        index = new_item.args[0]
        if index.has(n):
            m_indexs.append(index)
    m_indexs = set(m_indexs)
    m_indexs = list(m_indexs)
    max_index = None
    min_index = None
    for item in m_indexs:
        if item.free_symbols:
            max_index = item
            min_index = item
            break
    for item in m_indexs:
        max_minus = item - max_index
        if not max_minus.free_symbols:
            if max_minus >= 0:
                max_index = item
        min_minus = item - min_index
        if not min_minus.free_symbols:
            if min_minus < 0:
                min_index = item
    return max_index, min_index


# 求递推公式下标的最值
def sl_max_min_index002(poly):
    p = r'[a-zA-Z]_\(.*?\)'
    m = re.findall(p, str(poly))
    m = list(set(m))
    m_indexs = []
    for item in m:
        new_item = simplify(item)
        index = new_item.args[0]
        m_indexs.append(index)
    m_indexs = set(m_indexs)
    m_indexs = list(m_indexs)
    max_index = m_indexs[0]
    min_index = m_indexs[0]
    for item in m_indexs:
        max_minus = item - max_index
        if max_minus >= 0:
            max_index = item
        min_minus = item - min_index
        if min_minus < 0:
            min_index = item
    return max_index, min_index


# 求递推公式的数列名
def sl_shulie_name(poly):
    p = r'[a-zA-Z]_\(.*?\)'
    m = re.findall(p, str(poly))
    m = list(set(m))
    shulie_name = []
    for item in m:
        new_item = sympify(item)
        shulie_name.append(new_item.func)
    shulie_name = list(set(shulie_name))
    return shulie_name


# 把递推公式统一变成a(n+1),a(n)
def sl_new_an(poly):
    max_index, min_index = sl_max_min_index(poly)
    item = sympify('tt')
    answer = solve(min_index - item, n)[0]
    new_poly = poly.subs({n: answer})
    new_poly = new_poly.subs({item: n})
    return new_poly


# 判断an, sn是否是线性关系，如：S(n)+2a(n)=3*S(n-1)
def sl_sums001(poly):
    answer = False
    fn = simplify(poly)
    shulie_names = sl_shulie_name(fn)
    symbol = sl_default_symbol(fn)
    if len(shulie_names) == 2:
        if fn.is_Add:
            for arg in fn.args:
                xiangs = sl_xiangs(arg)
                if len(xiangs) == 2:
                    n_index = []
                    for item in xiangs:
                        item = sympify(item)
                        index = item.args[0]
                        if index.has(symbol):
                            n_index.append(1)
                        else:
                            n_index.append(0)
                    if sum(n_index) > 1:
                        answer = False
                        return answer
                    else:
                        answer = True
                else:
                    answer = True
        else:
            xiangs = sl_xiangs(fn)
            if len(xiangs) == 2:
                n_index = []
                for item in xiangs:
                    item = sympify(item)
                    index = item.args[0]
                    if index.has(symbol):
                        n_index.append(1)
                    else:
                        n_index.append(0)
                if sum(n_index) > 1:
                    answer = False
                    return answer
                else:
                    answer = True
    return answer


# 判断an, sn是否同时含Sn, an，如：S(n)+2a(n)=3*S(n-1)
def sl_sums002(poly):
    answer = False
    fn = simplify(poly)
    shulie_names = sl_shulie_name(fn)
    if len(shulie_names) == 2:
        answer = True
    return answer


# an, sn分开，如：S(n)+2a(n)=3*S(n-1)
def sl_apart(poly):
    fn = sympify(poly)
    shulie_names = sl_shulie_name(fn)
    symbol = sl_default_symbol(fn)
    s_ = sympify('S_')
    spart = None
    otherpart = None
    if len(shulie_names) == 2:
        if fn.is_Add:
            spart = 0
            for arg in fn.args:
                xiangs = sl_xiangs(arg)
                for item in xiangs:
                    item = sympify(item)
                    name = item.func
                    index = item.args[0]
                    if str(name) == str(s_) and index.has(symbol):
                        spart += arg
            otherpart = sympify(fn - spart)
    else:
        shulie_name = shulie_names[0]
        if str(shulie_name) == str("S_"):
            spart = fn
            otherpart = 0
        else:
            spart = 0
            otherpart = fn
    return spart, otherpart


def sl_simplify_sn(poly):
    fn = poly.simplify()
    p1 = r'[A-Z]_\(.*?\)'
    m1 = re.findall(p1, str(fn))
    m1 = list(set(m1))
    a_name = sympify("a_")
    if m1:
        for item in m1:
            new_item = sympify(item)
            new_item_index = new_item.args[0]
            new_item_value = 0
            if not new_item_index.free_symbols:
                for i in range(1, new_item_index + 1):
                    new_item_value += a_name(i)
                    fn = fn.subs({new_item: new_item_value})
    fn = fn.simplify()
    p2 = r'[A-Z]_\(.*?\)'
    m2 = re.findall(p2, str(fn))
    m2 = list(set(m2))
    s_inex = []
    for item in m2:
        new_item = sympify(item)
        new_item_index = new_item.args[0]
        s_inex.append(new_item_index)
    min_inex = s_inex[0]
    for item in s_inex:
        minus = (min_inex - item).expand().simplify()
        if not minus.free_symbols:
            if minus > 0:
                min_inex = item
    for item in m2:
        new_item = sympify(item)
        index = new_item.args[0]
        if index == min_inex:
            new_item_value = a_name(min_inex)
        else:
            minus = index - min_inex
            new_item_value = 0
            for i in range(0, minus + 1):
                new_item_value += a_name(min_inex + i)
        fn = fn.subs({new_item: new_item_value})
    return fn


# 归一化
def sl_gui_yi_hua_sa(eq):
    fn = eq[0] - eq[1]
    fn = sl_new_an(fn)
    spart1, otherpart1 = sl_apart(fn)
    eq_left = spart1
    eq_right = - otherpart1
    new_fn = fn.subs({n: n + 1})
    spart2, otherpart2 = sl_apart(new_fn)
    new_eq_left = spart2
    new_eq_right = - otherpart2
    left_minus = new_eq_left - eq_left
    new_left_minus = sl_simplify_sn(left_minus)
    right_minus = new_eq_right - eq_right
    return [new_left_minus, right_minus]


if __name__ == '__main__':
    pass
